#!/usr/bin/env python3
import rospy
import cv2
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from ultralytics import YOLO
import numpy as np
import torch
class YoloRosDetector:
    def __init__(self):
        rospy.init_node('yolo_ros_detector', anonymous=True)
        
        # 初始化模型
        self.model = YOLO("best.pt")  # 路径改成你的best.pt实际位置
        self.model.to("cuda")
        # 用于ROS图像转换
        self.bridge = CvBridge()
        print("是否启用CUDA:", torch.cuda.is_available())
        # 订阅图像话题
        self.image_sub = rospy.Subscriber("/camera/color/image_raw", Image, self.image_callback)

    def image_callback(self, msg):
        try:
            # 将ROS图像消息转为OpenCV格式（BGR）
            frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            rospy.logerr("图像转换失败: %s", str(e))
            return

        # 用YOLO进行推理
        results = self.model.predict(source=frame, imgsz=640, conf=0.3)

        # 在图像上画出检测框
        annotated_frame = results[0].plot()

        # 显示图像
        cv2.imshow("YOLOv11 Detection", annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            rospy.signal_shutdown("手动退出")

if __name__ == "__main__":
    try:
        detector = YoloRosDetector()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
    finally:
        cv2.destroyAllWindows()
